package io.github.novaframe.lock;

import org.apache.commons.lang3.StringUtils;
import org.springframework.dao.DataAccessException;
import org.springframework.data.redis.core.RedisCallback;
import org.springframework.data.redis.core.RedisOperations;
import org.springframework.data.redis.core.SessionCallback;
import org.springframework.data.redis.core.StringRedisTemplate;

import java.util.List;
import java.util.UUID;

/**
 * 使用Redis简单实现分布式锁机制
 *
 * @author Locki
 * @date Mar 31, 2017 6:02:56 PM
 * @version 1.0
 *
 */
public class RedisSimpleLock {

    final StringRedisTemplate redisTemplate;

    public RedisSimpleLock(StringRedisTemplate redisTemplate) {
        this.redisTemplate = redisTemplate;
    }

    /**
     * 获取永久锁，获取等待时间3秒
     * @param lockName
     * @return
     */
    public String getLock(String lockName) {
		long getTimeout = 3000;
	    // 随机128位UUID作为键的值
	    String identifier = UUID.randomUUID().toString();
	    final byte[] rawKey = redisTemplate.getStringSerializer().serialize("lock:" + lockName);
	    final byte[] rawVal = redisTemplate.getStringSerializer().serialize(identifier);

        long end = System.currentTimeMillis() + getTimeout;
	    while (System.currentTimeMillis() < end) {
	        RedisCallback<Boolean> redisCallback = redisConnection ->
	                /**
	                 * setnx命令的语义是将key的值设为value，当且仅当key不存在
	                 * 若key存在，不做任何动作，返回0(false)
	                 */
	                redisConnection.setNX(rawKey, rawVal);
	       if(redisTemplate.execute(redisCallback)){
	           return identifier;
	       }
	    }
	    return null;
	}

    /**
     * 获取超时锁，获取等待时间3秒
     * @param lockName
     * @param lockExpiredSec
     * @return
     */
	public String getLockExpired(String lockName, long lockExpiredSec) {
	    //3秒内获取不到锁就返回
	    long getTimeout = 3000;
        // 默认超时时间30秒
        if (lockExpiredSec <= 0) {
            lockExpiredSec = 30;
        }
	    String identifier = UUID.randomUUID().toString();
	    final byte[] rawKey = redisTemplate.getStringSerializer().serialize("lock:" + lockName);
	    final byte[] rawVal = redisTemplate.getStringSerializer().serialize(identifier);
	    long end = System.currentTimeMillis() + getTimeout;

	    while (System.currentTimeMillis() < end) {

            long finalLockExpiredSec = lockExpiredSec;
            RedisCallback<Boolean> redisCallback = redisConnection ->{
	            //成功获取锁之后设置锁超时时间
	            if(redisConnection.setNX(rawKey, rawVal)){
	                redisConnection.expire(rawKey, finalLockExpiredSec);
	                return true;
	            }
	            /**
	             * 锁获取失败之后检测锁是否有超时时间，如果没有则设置超时时间
	             * 这是为了防止程序在setNX和expire之间崩溃
	             */
	            if(redisConnection.ttl(rawKey) == -1){
	                redisConnection.expire(rawKey, finalLockExpiredSec);
	                return false;
	            }
	            return false;
	        };

	        if(redisTemplate.execute(redisCallback)){
	            return identifier;
	        }
	    }
	    return null;
	}

    /**
     * 释放锁
     * Redis cluster模式下使用，不支持redis事务
     * @param lockName
     * @param identifier
     * @return
     */
    public boolean releaseLock2(String lockName, String identifier) {
        String lockKey = "lock:" + lockName;
        String ret = redisTemplate.opsForValue().get(lockKey);

        if (StringUtils.equalsIgnoreCase(ret, identifier)) {
            redisTemplate.delete(lockKey);
            return true;
        }

        return false;
    }

    /**
     * 释放锁
     * Redis单机使用，支持Redis事务
     * @param lockName
     * @param identifier
     * @return
     */
	public boolean releaseLock(String lockName, String identifier) {
	    String lockKey = "lock:" + lockName;

	    SessionCallback<List<Object>> sessionCallback =  new SessionCallback<List<Object>>() {
	        @SuppressWarnings({ "rawtypes", "unchecked" })
			public List<Object> execute(RedisOperations operations) throws DataAccessException {
	            operations.watch(lockKey);
	            if(identifier.equals(operations.opsForValue().get(lockKey))){
	                operations.multi();
	                operations.delete(lockKey);
	                return operations.exec();
	            }

	            // 若取出来的不是想要释放的锁，不作任何操作，返回空
	            operations.unwatch();
	            return null;
	        }
	    };

	    List<Object> results = redisTemplate.execute(sessionCallback);
	    return results != null ? true : false;
	}
}
